#define _GNU_SOURCE
#include "threadpool.h"
#include "database.h"
#include "md5.h"
#include "scanner.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/stat.h>
#include <errno.h>
#include <time.h>

#define TASK_QUEUE_SIZE 1000
#define WAIT_TIMEOUT_SEC 5  // 超时检查间隔（秒）

typedef struct TaskNode {
    ThreadTask *task;
    struct TaskNode *next;
} TaskNode;

typedef struct {
    TaskNode *head;
    TaskNode *tail;
    int count;
    pthread_mutex_t mutex;
    pthread_cond_t not_empty;
    pthread_cond_t not_full;
} TaskQueue;

typedef struct {
    pthread_t thread_id;
    int active;
    ThreadPool *pool;
} WorkerThread;

struct ThreadPool {
    WorkerThread *threads;
    int thread_count;
    TaskQueue task_queue;
    int shutdown;
    int tasks_in_progress;
    pthread_mutex_t pool_mutex;
    pthread_cond_t all_tasks_done;
    ScannerContext *context;
};

// 任务队列初始化
static void task_queue_init(TaskQueue *queue) {
    queue->head = NULL;
    queue->tail = NULL;
    queue->count = 0;
    pthread_mutex_init(&queue->mutex, NULL);
    pthread_cond_init(&queue->not_empty, NULL);
    pthread_cond_init(&queue->not_full, NULL);
}

// 任务队列销毁
static void task_queue_destroy(TaskQueue *queue) {
    pthread_mutex_lock(&queue->mutex);
    TaskNode *current = queue->head;
    while (current) {
        TaskNode *next = current->next;
        if (current->task) {
            free(current->task);
        }
        free(current);
        current = next;
    }
    pthread_mutex_unlock(&queue->mutex);
    pthread_mutex_destroy(&queue->mutex);
    pthread_cond_destroy(&queue->not_empty);
    pthread_cond_destroy(&queue->not_full);
}

// 将任务加入队列
static int task_queue_push(TaskQueue *queue, ThreadTask *task, int shutdown) {
    pthread_mutex_lock(&queue->mutex);

    // 如果队列已满，等待
    while (queue->count >= TASK_QUEUE_SIZE && !shutdown) {
        pthread_cond_wait(&queue->not_full, &queue->mutex);
    }

    if (shutdown) {
        pthread_mutex_unlock(&queue->mutex);
        return -1;
    }

    TaskNode *node = malloc(sizeof(TaskNode));
    if (!node) {
        pthread_mutex_unlock(&queue->mutex);
        return -1;
    }

    node->task = task;
    node->next = NULL;

    if (queue->tail) {
        queue->tail->next = node;
    } else {
        queue->head = node;
    }
    queue->tail = node;
    queue->count++;

    pthread_cond_signal(&queue->not_empty);
    pthread_mutex_unlock(&queue->mutex);
    return 0;
}

// 从队列取出任务
static ThreadTask* task_queue_pop(TaskQueue *queue, int *shutdown) {
    pthread_mutex_lock(&queue->mutex);

    while (queue->count == 0 && !(*shutdown)) {
        pthread_cond_wait(&queue->not_empty, &queue->mutex);
    }

    if (*shutdown && queue->count == 0) {
        pthread_mutex_unlock(&queue->mutex);
        return NULL;
    }

    TaskNode *node = queue->head;
    ThreadTask *task = NULL;

    if (node) {
        task = node->task;
        queue->head = node->next;
        if (queue->head == NULL) {
            queue->tail = NULL;
        }
        queue->count--;
        free(node);
        pthread_cond_signal(&queue->not_full);
    }

    pthread_mutex_unlock(&queue->mutex);
    return task;
}

// 处理单个文件任务
static void process_file_task(ThreadTask *task, ScannerContext *context) {
    // 检查是否收到中断信号
    if (g_shutdown_requested || context->shutdown) {
        return;
    }
    
    struct stat st;
    if (stat(task->file_path, &st) != 0) {
        fprintf(stderr, "警告: 无法访问文件 %s: %s\n", task->file_path, strerror(errno));
        return;
    }

    // 只处理普通文件且大小符合要求的文件
    if (!S_ISREG(st.st_mode) || st.st_size < context->min_size) {
        return;
    }

    // 检查是否需要处理该文件（强制重算或数据库中不存在）
    int should_process = context->force_rehash;
    
    if (!should_process) {
        pthread_mutex_lock(&context->db_mutex);
        should_process = !db_file_exists(context->db, task->file_path);
        pthread_mutex_unlock(&context->db_mutex);
    }

    if (should_process) {
        FileInfo file_info;
        strncpy(file_info.file_name, task->file_path, MAX_PATH_LEN - 1);
        file_info.file_name[MAX_PATH_LEN - 1] = '\0';
        file_info.file_size = st.st_size;

        // 计算MD5
        int md5_result = calculate_md5(task->file_path, file_info.file_hash);
        if (md5_result == 0) {
            // 使用互斥锁保护数据库操作
            pthread_mutex_lock(&context->db_mutex);
            if (db_insert_or_update_file(context->db, &file_info) != SQLITE_OK) {
                fprintf(stderr, "错误: 无法插入文件信息到数据库 %s\n", task->file_path);
            }
            pthread_mutex_unlock(&context->db_mutex);
        } else if (md5_result == -2) {
            // MD5计算被中断，直接返回
            return;
        } else {
            fprintf(stderr, "错误: 无法计算文件MD5 %s\n", task->file_path);
        }
    }

    // 更新进度显示
    pthread_mutex_lock(&context->progress_mutex);
    context->processed_files++;
    printf("[%d/%d] %s\n", context->processed_files, context->total_files, task->file_path);
    pthread_mutex_unlock(&context->progress_mutex);
}

// 工作线程函数
static void* worker_thread_func(void *arg) {
    WorkerThread *worker = (WorkerThread *)arg;
    ThreadPool *pool = worker->pool;
    ScannerContext *context = pool->context;

    while (1) {
        // 检查是否收到中断信号
        if (g_shutdown_requested) {
            context->shutdown = 1;
            break;
        }

        // 从队列获取任务
        ThreadTask *task = task_queue_pop(&pool->task_queue, &pool->shutdown);

        if (task == NULL) {
            // 收到退出信号
            break;
        }

        // 再次检查中断信号，防止在获取任务后立即收到中断
        if (g_shutdown_requested) {
            context->shutdown = 1;
            free(task);
            break;
        }

        // 增加正在处理的任务计数
        pthread_mutex_lock(&pool->pool_mutex);
        pool->tasks_in_progress++;
        pthread_mutex_unlock(&pool->pool_mutex);

        // 处理任务
        process_file_task(task, context);
        free(task);

        // 减少正在处理的任务计数
        pthread_mutex_lock(&pool->pool_mutex);
        pool->tasks_in_progress--;
        if (pool->tasks_in_progress == 0 && pool->task_queue.count == 0) {
            pthread_cond_broadcast(&pool->all_tasks_done);
        }
        pthread_mutex_unlock(&pool->pool_mutex);
    }

    return NULL;
}

ThreadPool* threadpool_create(int thread_count, ScannerContext *context) {
    if (thread_count > MAX_THREAD_COUNT) {
        thread_count = MAX_THREAD_COUNT;
    }

    ThreadPool *pool = malloc(sizeof(ThreadPool));
    if (!pool) return NULL;

    pool->threads = malloc(sizeof(WorkerThread) * thread_count);
    if (!pool->threads) {
        free(pool);
        return NULL;
    }

    pool->thread_count = thread_count;
    pool->shutdown = 0;
    pool->tasks_in_progress = 0;
    pool->context = context;
    pthread_mutex_init(&pool->pool_mutex, NULL);
    pthread_cond_init(&pool->all_tasks_done, NULL);
    task_queue_init(&pool->task_queue);

    // 创建工作线程
    int created_threads = 0;
    for (int i = 0; i < thread_count; i++) {
        WorkerThread *worker = &pool->threads[i];
        worker->active = 0;
        worker->pool = pool;

        if (pthread_create(&worker->thread_id, NULL, worker_thread_func, worker) == 0) {
            worker->active = 1;
            created_threads++;
        } else {
            fprintf(stderr, "错误: 无法创建工作线程 %d\n", i);
        }
    }

    if (created_threads == 0) {
        task_queue_destroy(&pool->task_queue);
        pthread_mutex_destroy(&pool->pool_mutex);
        pthread_cond_destroy(&pool->all_tasks_done);
        free(pool->threads);
        free(pool);
        return NULL;
    }

    pool->thread_count = created_threads;
    return pool;
}

int threadpool_submit(ThreadPool *pool, ThreadTask *task) {
    if (!pool || !task || pool->shutdown) {
        return -1;
    }

    return task_queue_push(&pool->task_queue, task, pool->shutdown);
}

void threadpool_wait(ThreadPool *pool, ScannerContext *context) {
    if (!pool) return;

    pthread_mutex_lock(&pool->pool_mutex);
    
    while (pool->tasks_in_progress > 0 || pool->task_queue.count > 0) {
        // 检查是否收到关闭信号
        if (g_shutdown_requested || pool->shutdown || (context && context->shutdown)) {
            break;
        }

        // 使用带超时的等待，防止死锁
        struct timespec timeout;
        clock_gettime(CLOCK_REALTIME, &timeout);
        timeout.tv_sec += WAIT_TIMEOUT_SEC;

        int wait_result = pthread_cond_timedwait(&pool->all_tasks_done, &pool->pool_mutex, &timeout);

        if (wait_result == ETIMEDOUT) {
            // 超时后检查线程状态
            int active_count = 0;
            for (int i = 0; i < pool->thread_count; i++) {
                if (pool->threads[i].active) {
                    active_count++;
                }
            }

            // 如果没有活跃线程但还有任务，说明可能出现问题
            if (active_count == 0 && (pool->tasks_in_progress > 0 || pool->task_queue.count > 0)) {
                break;
            }

            // 继续等待
            continue;
        }
    }
    
    pthread_mutex_unlock(&pool->pool_mutex);

    if (!g_shutdown_requested && !pool->shutdown && (!context || !context->shutdown)) {
        printf("所有文件处理完成！\n");
    }
}

void threadpool_destroy(ThreadPool *pool) {
    if (!pool) return;

    // 设置关闭标志
    pthread_mutex_lock(&pool->pool_mutex);
    pool->shutdown = 1;
    pthread_mutex_unlock(&pool->pool_mutex);

    // 唤醒所有等待的线程
    pthread_mutex_lock(&pool->task_queue.mutex);
    pthread_cond_broadcast(&pool->task_queue.not_empty);
    pthread_cond_broadcast(&pool->task_queue.not_full);
    pthread_mutex_unlock(&pool->task_queue.mutex);

    // 统一使用超时等待策略，确保所有线程都能退出
    struct timespec ts;
    int max_wait_time = g_shutdown_requested ? 3 : 5;  // 中断时等待3秒，正常退出等待5秒

    for (int i = 0; i < pool->thread_count; i++) {
        WorkerThread *worker = &pool->threads[i];
        if (worker->active) {
            clock_gettime(CLOCK_REALTIME, &ts);
            ts.tv_sec += max_wait_time;

            int result = pthread_timedjoin_np(worker->thread_id, NULL, &ts);
            if (result == ETIMEDOUT) {
                pthread_cancel(worker->thread_id);
                pthread_join(worker->thread_id, NULL);
            }
        }
    }

    // 清理资源
    task_queue_destroy(&pool->task_queue);
    pthread_mutex_destroy(&pool->pool_mutex);
    pthread_cond_destroy(&pool->all_tasks_done);
    free(pool->threads);
    free(pool);
}
